package com.bff.gaia.runtime.maqy;

import com.esotericsoftware.minlog.Log;
import gurobi.GRBException;
import com.bff.gaia.api.common.functions.RichMapPartitionFunction;
import com.bff.gaia.api.java.tuple.Tuple2;
import com.bff.gaia.runtime.JobException;
import com.bff.gaia.runtime.maqy.util.LPUtil;
import com.bff.gaia.util.Collector;

import java.util.*;

/**
 * AttachPartition中统计的各个子任务中的数据量最终会汇集于此，按ip和hostname分组统计。
 * 然后计算任务分配比例并广播发送给PercentRangeBoundaryBuilder用于计算边界值。
 */
public class StatisticPartition<T> extends RichMapPartitionFunction<T, List<Integer>> {

	@Override
	public void mapPartition(Iterable<T> values, Collector<List<Integer>> out) throws Exception {
		List<Integer> proportions = new ArrayList<>();
		List<HostInfo> orderedHostInfos = new ArrayList<>();
		Map<String, Long> recordNumPerHostMap = new HashMap<>(); // [ip]@[hostname] —— recordNum，记录每个TM中的record个数
		Map<String, Integer> subTaskNumPerHostMap = new HashMap<>(); // [ip]@[hostname] —— subTaskNum，记录每个TM中运行的subTask个数

		// 根据subTaskIdx排序，保证得出的比值顺序和subTaskIdx的自然增序一致
		String logInfo = "";
		int i = 0;
		for (T record : values) {
			HostInfo hostInfo = (HostInfo) ((Tuple2) record).f0;
			Long count = (Long) ((Tuple2) record).f1;

			logInfo += String.format("%5s#%d %s: %d\n", "", ++i, hostInfo, count);

			orderedHostInfos.add(hostInfo);
			if (!recordNumPerHostMap.containsKey(hostInfo.value())) {
				recordNumPerHostMap.put(hostInfo.value(), count);
				subTaskNumPerHostMap.put(hostInfo.value(), 1);
			} else {
				recordNumPerHostMap.put(hostInfo.value(), recordNumPerHostMap.get(hostInfo.value()) + count);
				subTaskNumPerHostMap.put(hostInfo.value(), subTaskNumPerHostMap.get(hostInfo.value()) + 1);
			}
		}
		orderedHostInfos.sort(new Comparator<HostInfo>() {
			@Override
			public int compare(HostInfo o1, HostInfo o2) {
				return o1.getSubTaskIdx() - o2.getSubTaskIdx();
			}
		});
		Log.info("hfy", "[StatisticPartition] Original data distribution:\n" + logInfo.substring(0, logInfo.length() - 1));

		// 应用统计结果，得到locationInfo的集合
		LocationInfo[] locationInfos = new LocationInfo[orderedHostInfos.size()];
		i = 0;
		logInfo = "";
		for (HostInfo hostInfo : orderedHostInfos) {
			long recordNum = recordNumPerHostMap.get(hostInfo.value());
			int sumTaskNum = subTaskNumPerHostMap.get(hostInfo.value());

			// 相同TM下的subTasks平分数据量是为了保证平均分配比例，因为运行在相同TM下的subTasks共享网络带宽
			long subRecordNum = recordNum / sumTaskNum; // 即便不能整除，误差也可以忽略不计
			locationInfos[i++] = new LocationInfo(hostInfo.getHostName(), hostInfo.getIp(), (int) hostInfo.getRate().getUpLink(), (int) hostInfo.getRate().getDownLink(), subRecordNum);

			logInfo += String.format("%5s#%d %s: %d\n", "", i, hostInfo, subRecordNum);
		}
		Log.info("hfy", "[StatisticPartition] Data distribution after rebalance:\n" + logInfo.substring(0, logInfo.length() - 1));

		// 计算proportions
		try {
			new LPUtil().getLocationProportion(locationInfos);
		} catch (GRBException e) {
			e.printStackTrace();
			throw new JobException("LP计算出现错误");
		}

		// 得到比例数组
		logInfo = "[StatisticPartition] Proportions sorted by subTaskIdx: ";
		for (int idx = 0; idx < locationInfos.length; idx++) {
			int proportion = locationInfos[idx].getProportion();

			logInfo += proportion + ":";

			if (proportion > 100 || proportion < 0) {
				//throw new Exception("");
				System.out.println("proportion is incorrect");
			}
			proportions.add(proportion);
		}
		logInfo = logInfo.substring(0, logInfo.length() - 1);
		Log.info("hfy", logInfo);

		out.collect(proportions);
	}
}